package translationModel;

import java.util.HashMap;
import java.util.PriorityQueue;
import java.util.Set;

public class StateStack
{
	private HashMap<State, State> stack;
	private int beamSize;
	
	public StateStack()
	{
		this(-1);
	}
	
	public StateStack(int beamSize)
	{
		stack = new HashMap<State, State>();
		this.beamSize = beamSize;
	}
	
	public Set<State> GetStates()
	{
		return stack.keySet();
	}
	
	public void AddState(State state)
	{
		// Check if recombination needed
		State stateInStack = stack.get(state);
		if (stateInStack == null)
			stack.put(state, state);
		else if (stateInStack.score < state.score)
		{
			// Recombine the given state with an existing state
			stateInStack.score = state.score;
			stateInStack.prevState = state.prevState;
			stateInStack.phrase = state.phrase;
		}
	}
	
	public void Beam()
	{
		// Check if beam needed
		if ((beamSize == -1) || (stack.size() < beamSize))
			return;
		// Sort the states using Priority Queue
		PriorityQueue<State> queue = new PriorityQueue<State>(stack.keySet());
		// Leave the best score states
		for (int i = 0; i < beamSize; i++)
			queue.poll();
		// Remove all the rest
		for (State state = queue.poll(); state != null; state = queue.poll())
			stack.remove(state);
	}
	
	public State GetBestScore()
	{
		State best = null;
		for (State state : stack.keySet())
		{
			if ((best == null) || (state.score > best.score))
				best = state;
		}
		return best;
	}
	
	public State GetBestEEnd(int eEnd)
	{
		State best = null;
		for (State state : stack.keySet())
		{
			if ((state.eEndPoint == eEnd) && ((best == null) || (state.score > best.score)))
				best = state;
		}
		return best;
	}
}
